import cv2
import supervision as sv
from rfdetr import RFDETRBase
from collections import defaultdict
from typing import Dict, Set
from PIL import Image, ImageDraw, ImageFont  # 导入PIL库
import numpy as np  # 导入numpy用于图像格式转换

class RFDETRDetector:
    def __init__(self, model_path='models/checkpoint_best_ema.pth', resolution=560, font_path="C:/Windows/Fonts/simhei.ttf", font_size=15):
        # 1. 初始化模型
        self.model = RFDETRBase(
            pretrain_weights=model_path,
            # pretrain_weights=model_path or r"E:\A\rf-detr-main\output\pre-train1\checkpoint_best_ema.pth",
            resolution=resolution
        )

        # 2. 初始化跟踪器
        # TODO: 根据需要调整参数
        self.tracker = sv.ByteTrack(
            track_activation_threshold=0.5,
            lost_track_buffer=120,
            minimum_matching_threshold=0.9, # 注意：之前代码是0.95，本地是0.9，根据实际效果调整
            minimum_consecutive_frames=2,
            frame_rate=20 # 假设帧率，如果可能，应动态获取或配置
        )

        # 3. 类别定义
        self.VISDRONE_CLASSES = [
            'pedestrian', 'people', 'bicycle', 'car', 'van',
            'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor'
        ]
        self.VISDRONE_CLASSES_CHINESE = {
            'pedestrian': '行人', 'people': '人', 'bicycle': '自行车', 'car': '小汽车', 'van': '面包车',
            'truck': '卡车', 'tricycle': '三轮车', 'awning-tricycle': '带篷三轮车', 'bus': '公交车', 'motor': '摩托车'
        }

        # 4. 初始化字体
        self.FONT_SIZE = font_size
        try:
            self.font = ImageFont.truetype(font_path, self.FONT_SIZE)
        except IOError:
            print(f"错误：无法加载字体 {font_path}。请确保路径正确且文件存在。")
            self.font = ImageFont.load_default()

        # 5. 类别计数器 (作为类属性)
        self.class_tracks: Dict[str, Set[int]] = defaultdict(set)
        self.category_counts: Dict[str, int] = defaultdict(int)

        # 6. 初始化标注器 (仅保留 BoxAnnotator)
        self.box_annotator = sv.BoundingBoxAnnotator(
            thickness=2,
            color=sv.Color(r=0, g=255, b=0) # 绿色边框
        )

    def _update_counter(self, detections: sv.Detections):
        """更新类别计数器"""
        # 只统计有 tracker_id 的检测结果
        valid_indices = detections.tracker_id != None
        if not np.any(valid_indices): # 处理 detections 为空或 tracker_id 都为 None 的情况
             return

        class_ids = detections.class_id[valid_indices]
        track_ids = detections.tracker_id[valid_indices]

        for class_id, track_id in zip(class_ids, track_ids):
            if track_id is None: # 跳过没有 tracker_id 的项
                continue
            # 使用英文类别名作为内部 key
            class_name = self.VISDRONE_CLASSES[class_id]
            if track_id not in self.class_tracks[class_name]:
                self.class_tracks[class_name].add(track_id)
                self.category_counts[class_name] += 1

    def _draw_frame(self, frame: np.ndarray, detections: sv.Detections) -> np.ndarray:
        """绘制检测框、中文标签和计数信息"""
        # 绘制检测框
        annotated_frame = self.box_annotator.annotate(scene=frame.copy(), detections=detections)

        # --- 使用 PIL 绘制中文标签和统计面板 ---
        pil_image = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)

        # 绘制中文标签
        valid_indices = detections.tracker_id != None
        if np.any(valid_indices): # 检查是否有有效的检测结果
            boxes = detections.xyxy[valid_indices]
            class_ids = detections.class_id[valid_indices]
            track_ids = detections.tracker_id[valid_indices] # 虽然下面没用到 track_id，但保持一致性

            for box, class_id in zip(boxes, class_ids):
                x1, y1, _, _ = map(int, box) # 使用 int() 转换坐标

                # 获取英文类别名，再通过映射获取中文类别名
                english_label = self.VISDRONE_CLASSES[class_id]
                chinese_label = self.VISDRONE_CLASSES_CHINESE.get(english_label, english_label)

                text_to_draw = f"{chinese_label}"
                text_color = (255, 255, 255) # 白色 (RGB)

                # 计算文本位置
                text_x = x1
                text_y = y1 - self.FONT_SIZE - 2
                if text_y < 0:
                    text_y = y1 + 2

                draw.text((text_x, text_y), text_to_draw, font=self.font, fill=text_color)

        # 绘制统计面板 (右上角)
        stats_text_lines = [
            f"{self.VISDRONE_CLASSES_CHINESE.get(cls, cls)}: {self.category_counts[cls]}"
            for cls in self.VISDRONE_CLASSES if self.category_counts[cls] > 0
        ]

        frame_height, frame_width, _ = frame.shape # 从输入帧获取宽高
        stats_start_x = frame_width - 200 # 动态计算位置
        stats_start_y = 10
        line_height = self.FONT_SIZE + 5
        stats_text_color = (255, 255, 255)

        for i, line in enumerate(stats_text_lines):
            text_pos = (stats_start_x, stats_start_y + i * line_height)
            draw.text(text_pos, line, font=self.font, fill=stats_text_color)

        # --- 所有 PIL 绘制完成后，统一转换回 OpenCV 格式 ---
        final_annotated_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        return final_annotated_frame

    def detect_and_draw_count(self, frame: np.ndarray, conf: float = 0.3) -> np.ndarray:
        """执行单帧检测、跟踪、计数并绘制结果"""
        try:
            # 1. 执行检测
            detections = self.model.predict(frame, threshold=conf)

            # 处理 detections 为 None 或空的情况
            if detections is None or len(detections) == 0:
                # 如果没有检测结果，仍然需要绘制统计面板（如果之前有计数）
                # 或者直接返回原始帧加上可能的统计面板
                # 为简单起见，我们可以在 _draw_frame 中处理空的 detections
                 # 创建一个空的 Detections 对象以便绘制统计信息
                 detections = sv.Detections.empty() # 需要确保 _draw_frame 能处理空 detections
                 # 注意：空的 detections 没有 tracker_id，因此 _update_counter 不会执行
                 # 如果需要绘制旧的统计信息，需要确保 _draw_frame 能访问 self.category_counts
                 annotated_frame = self._draw_frame(frame, detections) # 绘制统计面板
                 return annotated_frame


            # 2. 执行跟踪
            detections = self.tracker.update_with_detections(detections)

            # 3. 更新计数器
            self._update_counter(detections)

            # 4. 绘制结果
            annotated_frame = self._draw_frame(frame, detections)

            return annotated_frame

        except Exception as e:
            print(f"处理帧时发生错误: {e}")
            # 发生错误时返回原始帧，避免程序中断
            return frame